#!/usr/bin/env python
# coding: utf-8

# # 保存与加载
# 上一章节的内容里面主要是介绍了如何调整超参数，并进行网络模型训练。训练网络模型的过程中，实际上我们希望保存中间和最后的结果，用于微调（fine-tune）和后续的模型部署和推理，本章节我们开始学习如何保存与加载模型。
#
# ## 模型训练
#
# 下面我们以MNIST数据集为例，介绍网络模型的保存与加载方式。首先，我们需要获取MNIST数据集并训练模型，示例代码如下：

# In[2]:


import luojianet.nn as nn
import luojianet as ms

#from mindvision.classification.dataset import Mnist
from luojianet.dataset import MnistDataset
from src.lenet import lenet
from luojianet.train.callback import LossMonitor
from luojianet.dataset.vision.c_transforms import RandomHorizontalFlip, HWC2CHW
from luojianet.dataset.transforms.c_transforms import TypeCast
from luojianet.common import dtype as mstype
epochs = 1  # 训练轮次

# 1. 构建数据集
dataset_train=MnistDataset(dataset_dir='./MNIST_data',usage='train')
# 数据增强
trans = [
    RandomHorizontalFlip(prob=0.5),
    HWC2CHW(),
    TypeCast(mstype.float32)
]

convert = [
    TypeCast(mstype.int32)
]

dataset_train = dataset_train.map(operations=trans, input_columns=['image'])
dataset_train = dataset_train.map(operations=convert, input_columns=['label']).batch(32, True)
# 2. 定义神经网络
network = lenet()
# 3.1 定义损失函数
net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
# 3.2 定义优化器函数
net_opt = nn.Momentum(network.trainable_params(), learning_rate=0.01, momentum=0.9)
# 3.3 初始化模型参数
model = ms.Model(network, loss_fn=net_loss, optimizer=net_opt, metrics={'accuracy'})

# 4. 对神经网络执行训练
model.train(epochs, dataset_train, callbacks=LossMonitor())


# 从上面的打印结果可以看出，随着训练轮次的增加，损失值趋于收敛。
#
# ## 保存模型
#
# 在训练完网络完成后，下面我们将网络模型以文件的形式保存下来。保存模型的接口有主要2种：
#
# 1. 简单的对网络模型进行保存，可以在训练前后进行保存。这种方式的优点是接口简单易用，但是只保留执行命令时候的网络模型状态；
# 2. 在网络模型训练中进行保存，luojianet在网络模型训练的过程中，自动保存训练时候设定好的epoch数和step数的参数，也就是把模型训练过程中产生的中间权重参数也保存下来，方便进行网络微调和停止训练；
#
# ### 直接保存模型
#
# 使用luojianet提供的save_checkpoint保存模型，传入网络和保存路径：

# In[3]:


import luojianet as ms

# 定义的网络模型为net，一般在训练前或者训练后使用
ms.save_checkpoint(network, "MyNet.ckpt")
print("saved MyNet.ckpt")

# 其中，`network`为训练网络，`"./MyNet.ckpt"`为网络模型的保存路径。
#
# ### 训练过程中保存模型
#
# 在模型训练的过程中，使用`model.train`里面的`callbacks`参数传入保存模型的对象 [ModelCheckpoint](https://luojianet.cn/docs/zh-CN/r1.8/api_python/luojianet/luojianet.ModelCheckpoint.html#luojianet.ModelCheckpoint)（一般与[CheckpointConfig](https://luojianet.cn/docs/zh-CN/r1.8/api_python/luojianet/luojianet.CheckpointConfig.html#luojianet.CheckpointConfig)配合使用），可以保存模型参数，生成CheckPoint(简称ckpt)文件。
#
# 用户可以根据具体需求通过设置`CheckpointConfig`来对CheckPoint策略进行配置。具体用法如下：

# In[4]:



# 设置epoch_num数量
epoch_num = 5

from luojianet.train.callback import ModelCheckpoint, CheckpointConfig
# 设置模型保存参数
config_ck = CheckpointConfig(save_checkpoint_steps=100, keep_checkpoint_max=10)

# 应用模型保存参数
# 每100个step(100个batch批次)保存一次模型参数，最多保存10个
ckpoint = ModelCheckpoint(prefix="lenet", directory="./lenet", config=config_ck)
model.train(epoch_num, dataset_train, callbacks=[ckpoint])
print("saved lenet.ckpt")

# 上述代码中，首先需要初始化一个`CheckpointConfig`类对象，用来设置保存策略。
#
# - `save_checkpoint_steps`表示每隔多少个step保存一次。
# - `keep_checkpoint_max`表示最多保留CheckPoint文件的数量。
# - `prefix`表示生成CheckPoint文件的前缀名。
# - `directory`表示存放文件的目录。
#
# 创建一个`ModelCheckpoint`对象把它传递给`model.train`方法，就可以在训练过程中使用CheckPoint功能了。
#
# 生成的CheckPoint文件如下：
#
# ```text
# lenet-graph.meta # 编译后的计算图
# lenet-1_1875.ckpt  # CheckPoint文件后缀名为'.ckpt'
# lenet-2_1875.ckpt  # 文件的命名方式表示保存参数所在的epoch和step数，这里为第2个epoch的第1875个step的模型参数
# lenet-3_1875.ckpt  # 表示保存的是第3个epoch的第1875个step的模型参数
# ...
# ```
#
# 如果用户使用相同的前缀名，运行多次训练脚本，可能会生成同名CheckPoint文件。luojianet为方便用户区分每次生成的文件，会在用户定义的前缀后添加"_"和数字加以区分。如果想要删除`.ckpt`文件时，请同步删除`.meta` 文件。
#
# 例：`lenet_3-2_1875.ckpt` 表示运行第3次脚本生成的第2个epoch的第1875个step的CheckPoint文件。
#
# ## 加载模型
#
# 要加载模型权重，需要先创建相同模型的实例，然后使用`load_checkpoint`和`load_param_into_net`方法加载参数。
#
# 示例代码如下：

# In[5]:


import luojianet as ms

from luojianet.dataset import MnistDataset
from src.lenet import lenet

# 将模型参数存入parameter的字典中，这里加载的是上面训练过程中保存的模型参数
param_dict = ms.load_checkpoint("./lenet/lenet-5_1875.ckpt")

# 重新定义一个LeNet神经网络
net = lenet()

# 将参数加载到网络中
ms.load_param_into_net(net, param_dict)

# 重新定义优化器函数
net_opt = nn.Momentum(net.trainable_params(), learning_rate=0.01, momentum=0.9)

model = ms.Model(net, loss_fn=net_loss, optimizer=net_opt, metrics={"accuracy"})


# - `load_checkpoint`方法会把参数文件中的网络参数加载到字典`param_dict`中。
# - `load_param_into_net`方法会把字典`param_dict`中的参数加载到网络或者优化器中，加载后，网络中的参数就是CheckPoint保存的。
#
# ### 模型验证
#
# 在上述模块把参数加载到网络中之后，针对推理场景，可以调用`eval`函数进行推理验证。示例代码如下：

# In[8]:


# 调用eval()进行推理
#download_eval = Mnist(path="./mnist", split="test", batch_size=32, resize=32, download=True)
#dataset_eval = download_eval.run()
dataset_eval=MnistDataset(dataset_dir='./MNIST_data',usage='test')

dataset_eval = dataset_eval.map(operations=trans, input_columns=['image'])
dataset_eval = dataset_eval.map(operations=convert, input_columns=['label']).batch(32, True)

acc = model.eval(dataset_eval)

print("{}".format(acc))


# ### 用于迁移学习
#
# 针对任务中断再训练及微调（Fine-tuning）场景，可以调用`train`函数进行迁移学习。示例代码如下：

# In[9]:


# 定义训练数据集
#download_train = Mnist(path="./mnist", split="train", batch_size=32, repeat_num=1, shuffle=True, resize=32, download=True)
#dataset_train = download_train.run()
dataset_train=MnistDataset(dataset_dir='./MNIST_data',usage='train')
dataset_train = dataset_train.map(operations=trans, input_columns=['image'])
dataset_train = dataset_train.map(operations=convert, input_columns=['label']).batch(32, True)
# 网络模型调用train()继续进行训练
model.train(1, dataset_train, callbacks=LossMonitor())
